﻿using System;
using System.Collections.Generic;
using Allegro.Framework;
//using Cosmologica.Framework;

namespace Allegro.Integrator
{
    public class ButcherTableau
    {
        public int nStages;
        public double[] c; // x interpolation coeff.
        public double[,] a; // y interpolation coeff.
        public double[] b; // High-order weights
        public double[] bh; // Low order weights
        public double[] db; // db = b - bh

        public ButcherTableau(int n)
        {
            nStages = n;
            c = new double[n];
            a = new double[n, n];
            b = new double[n];
            bh = new double[n];
            db = new double[n];
        }
    }
    /// <summary>
    /// Generic implementation of the Runge-Kutta integration. The specific method
    /// is given by specifying the Butcher tableau.
    /// </summary>
    public class GenericRungeKuttaIntegrator : IODEIntegrator
    {
        public const double TINY = 1.0e-30;

        /// <summary>
        /// Default accuracy
        /// </summary>
        public const double DefaultAccuracy = 1.0e-8;
        const double SAFETY = 0.9;
        const double PGROW = -0.2;
        const double PSHRNK = -0.25;
        const double ERRCON = 1.89e-4; // The value ERRCON equals (5/SAFETY) raised to the power (1/PGROW), see use below
        const double MaxStepReductionFactor = 0.1;
        const double MaxStepIncreaseFactor = 5;

        private double[] c;
        private double[,] a;
        private double[] b;
        private double[] bh;
        private double[] db;

        private ButcherTableau _tableau;
        int nStages; // Number of stages
        int _n;
        private double[][] _ak;
        double[] _ytemp;
        double[] _yComputed;
        double[] _yerr;
        double[] _dydx;
        double[] _yscal;

        double _accuracy = DefaultAccuracy;

        bool[] _integrated; // Flag for each parameter. If true, variable i was integrated.
        bool[] _computed; // Flag for each parameter. If true, variable i was computed directly

        /// <summary>
        /// 
        /// </summary>
        /// <param name="dimension"></param>
        public GenericRungeKuttaIntegrator(int dimension, ButcherTableau tableau)
        {
            if (dimension <= 0)
                throw new ApplicationException("Dim. must be > 0");

            nStages = tableau.nStages;

            _tableau = tableau;
            _n = dimension;
            _ak = new double[nStages][];
            for (int i = 0; i < nStages; i++)
            {
                _ak[i] = new double[_n];
            }
            int n = _tableau.nStages;
            c = _tableau.c;
            a = _tableau.a;
            b = _tableau.b;
            bh = _tableau.bh;
            db = _tableau.db;

            // Error weights
            for (int i = 0; i < n; i++)
            {
                db[i] = b[i] - bh[i];
            }

            _ytemp = new double[_n];
            _yComputed = new double[_n];
            _yerr = new double[_n];
            _dydx = new double[_n];
            _yscal = new double[_n];

            _integrated = new bool[_n];
            _computed = new bool[_n];
            _ResetCompletedFlags();
        }

        /// <summary>
        /// Relative accuracy
        /// </summary>
        public double Accuracy
        {
            get
            {
                return _accuracy;
            }
            set
            {
                _accuracy = value;
            }
        }

        /// <summary>
        /// Integration step. Input are the dependent variable vector y[1..n] and its derivative dydx[1..n]
        /// at the starting value of the independent varable x. Also input are the stepsize to be attempted, 
        /// htry, and the vector yscal[1..n], against which the error is scaled. On output, y and x are replaced by their
        /// new values, hdid is the stepsize that was actually accomplished, and hnext is the estimated next stepsize. 
        /// derivs is the user-supplied method that computes the right-hand side derivatives
        /// </summary>        
        /// <param name="x"></param>
        /// <param name="y"></param>
        /// <param name="htry"></param>
        /// <param name="hdid"></param>
        /// <param name="hnext"></param>
        /// <param name="derivs"></param>
        /// <param name="values"></param>
        public void Integrate(ref double x, ref double[] y, double htry, out double hdid, out double hnext,
            Derivative derivs, ComputeValues values)
        {
            double errmax, h, htemp, xnew;
            double eps = Accuracy;

            _ResetCompletedFlags();
            bool isComplete = derivs(x, y, ref _dydx, ref _integrated);

            h = htry; // Set stepsize to initial trial value
            for (int i = 0; i < _n; i++)
            {
                // Scaling used to monitor accuracy. This general-purpose choice can be modified if need be.
                _yscal[i] = Math.Abs(y[i]) + Math.Abs(_dydx[i] * h) + TINY;
            }

            //double[] errors = new double[_n];
            //int iErrorMax = 0;
            for (; ; )
            {
                // Take a step
                _MakeIntegrationStep(y, _dydx, x, h, ref _ytemp, ref _yerr, derivs, values);
                errmax = 0.0;
                for (int i = 0; i < _n; i++)
                {
                    //errors[i] = Math.Abs(_yerr[i] / _yscal[i])/eps;
                    //errors[i] = Math.Abs(_yerr[i]) / eps;
                    //if (errors[i] > errmax / eps)
                    //    iErrorMax = i;
                    //errmax = Math.Max(errmax, Math.Abs(_yerr[i] / _yscal[i]));
                    errmax = Math.Max(errmax, Math.Abs(_yerr[i]));
                }
                errmax /= eps;
                if (errmax <= 1.0)
                    break;
                htemp = SAFETY * h * Math.Pow(errmax, PSHRNK);
                // Truncation error too large, reduce stepsize
                // Limit the step reduction by MaxStepReductionFactor (0.1 by default)
                h = (h > 0.0 ? Math.Max(htemp, MaxStepReductionFactor * h) : Math.Min(htemp, MaxStepReductionFactor * h));
                xnew = x + h;
                if (xnew == x)
                    ErrorHandler.Error(string.Format("Stepsize underflow in RungeKuttaIntegrator.Integrate occurred at x={0}", x));
            }
            if (errmax > ERRCON)
                hnext = SAFETY * h * Math.Pow(errmax, PGROW);
            else
                hnext = MaxStepIncreaseFactor * h; // No more than a factor of 5 increase
            x += h;
            hdid = h;
            for (int i = 0; i < _n; i++)
            {
                y[i] = _ytemp[i];
            }
        }

        // Private methods
        void _ResetCompletedFlags()
        {
            for (int i = 0; i < _n; i++)
            {
                _computed[i] = _integrated[i] = false;
            }
        }
        /// <summary>
        /// Given values for n variables y[1..n] and their derivatives dydx[1..n] known at x, use 
        /// the Runge-Kutta method to advance the solution over an interval h and
        /// return the incremented variables as yout[1..n]. Also return an estimate of the local
        /// truncation error in yout using the embedded fourth-order method. The method invokes the 
        /// user-supplied method derivs through the delegat Derivative(x,y,dydx), 
        /// which computes the first-order derivatives at x.
        /// The method accomodates any Runge-Kutta method that can be described via a Butcher tableau.
        /// </summary>
        /// <param name="y"></param>
        /// <param name="dydx"></param>
        /// <param name="x"></param>
        /// <param name="h"></param>
        /// <param name="yout"></param>
        /// <param name="yerr"></param>
        /// <param name="derivs"></param>
        void _MakeIntegrationStep(double[] y, double[] dydx, double x, double h, ref double[] yout, ref double[] yerr, Derivative derivs, ComputeValues values)
        {
            double xtemp;
            bool isComplete = false;

            // First stage
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                    _ytemp[i] = y[i] + h * a[1, 0] * dydx[i];
            }
            if (!isComplete && values != null)
                isComplete = values(x, y, ref _yComputed, _integrated, ref _computed);
            for (int i = 0; i < _n; i++)
            {
                if (!_integrated[i] && _computed[i])
                    _ytemp[i] = _yComputed[i];
            }
            // Next stages
            double fAvg;
            for (int j = 1; j < nStages - 1; j++)
            {
                _ResetCompletedFlags();
                xtemp = x + c[j] * h;
                isComplete = derivs(xtemp, _ytemp, ref _ak[j], ref _integrated);
                for (int i = 0; i < _n; i++)
                {
                    fAvg = a[j + 1, 0] * dydx[i];
                    for (int k = 1; k <= j; k++)
                    {
                        fAvg += a[j + 1, k] * _ak[k][i];
                    }
                    if (_integrated[i])
                        _ytemp[i] = y[i] + h * fAvg;
                }
                if (!isComplete && values != null)
                    isComplete = values(xtemp, _ytemp, ref _yComputed, _integrated, ref _computed);
                for (int i = 0; i < _n; i++)
                {
                    if (!_integrated[i] && _computed[i])
                        _ytemp[i] = _yComputed[i];
                }
            }

            // Final stage
            int n = nStages - 1;
            _ResetCompletedFlags();
            xtemp = x + c[n] * h;
            isComplete = derivs(xtemp, _ytemp, ref _ak[n], ref _integrated);
            // Accumulate increments with proper weights
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                {
                    fAvg = b[0] * dydx[i];
                    for (int k = 1; k < nStages; k++)
                    {
                        fAvg += b[k] * _ak[k][i];
                    }
                    yout[i] = y[i] + h * fAvg;
                }
            }
            if (!isComplete && values != null)
                isComplete = values(xtemp, _ytemp, ref yout, _integrated, ref _computed);

            // Estimate error as difference between high and low order methods
            for (int i = 0; i < _n; i++)
            {
                if (_integrated[i])
                {
                    fAvg = db[0] * dydx[i];
                    for (int k = 1; k < nStages; k++)
                    {
                        fAvg += db[k] * _ak[k][i];
                    }
                    yerr[i] = h * fAvg / _yscal[i];
                }
                else
                    yerr[i] = 0.0;
            }
        }
    }
}
